#include "hook.h"

#include "module_enumerator.h"

#include <DbgHelp.h>

#include <algorithm>
#include <cstdint>
#include <type_traits>
#include <vector>

#pragma comment(lib, "Dbghelp.lib")

namespace windows
{

using namespace std::literals::string_view_literals;

static HMODULE get_module(const void* address) noexcept
{
    MEMORY_BASIC_INFORMATION info;
    return ::VirtualQuery(address, &info, sizeof(info)) ? reinterpret_cast<HMODULE>(info.AllocationBase) : NULL;
}

hook::hook(const std::vector<module_import>& import_functions) : current_module(get_module(this))
{
    for (const auto& module_functions : import_functions)
    {
        const auto module_handle = ::GetModuleHandleA(module_functions.module_name.data());

        if (module_handle)
        {
            std::string module_name = std::string(module_functions.module_name);
            std::transform(module_name.begin(), module_name.end(), module_name.begin(), ::tolower);

            auto& table = replaceable_tables[module_name];
            for (const auto& function : module_functions.replaceable_functions)
                table.emplace(std::string(function.first), hook::import_address{::GetProcAddress(module_handle, function.first.data()), function.second});

            module_names[module_handle] = std::move(module_name);
        }
    }

    module_enumerator enumerator;
    for (auto module_entry = enumerator.first(); module_entry; module_entry = enumerator.next())
        if (module_entry->hModule != current_module)
            for (const auto& item : import_functions)
                modify_addresses_in_import_tables(module_entry->hModule, operation_type::modify);
}

hook::~hook()
{
    module_enumerator enumerator;
    for (auto module_entry = enumerator.first(); module_entry; module_entry = enumerator.next())
        if (module_entry->hModule != current_module)
            modify_addresses_in_import_tables(module_entry->hModule, operation_type::cancel);
}

void hook::process_newly_module(const HMODULE module_address, const DWORD flags) noexcept
{
    if (!module_address || (flags & LOAD_LIBRARY_AS_DATAFILE))
        return;

    hook::get_or_create()->modify_addresses_in_import_tables(module_address, hook::operation_type::modify);
}

std::unique_ptr<hook>& hook::get_or_create(std::vector<module_import> import_functions)
{
    static const auto functions_hook = [&import_functions]()
    {
        import_functions.push_back({kernel_module, {{"LoadLibraryA"sv, reinterpret_cast<PROC>(load_library<PCSTR>)},
                                                       {"LoadLibraryW"sv, reinterpret_cast<PROC>(load_library<PCWSTR>)},
                                                       {"LoadLibraryExA"sv, reinterpret_cast<PROC>(load_library_ex<PCSTR>)},
                                                       {"LoadLibraryExW"sv, reinterpret_cast<PROC>(load_library_ex<PCWSTR>)},
                                                       {"GetProcAddress"sv, reinterpret_cast<PROC>(get_procedure_address)}}});

        return std::make_unique<hook>(import_functions);
    }();

    return functions_hook;
}

hook::import_address hook::get_function_address(std::string module_name, const std::string_view function_name)
{
    std::transform(module_name.begin(), module_name.end(), module_name.begin(), ::tolower);

    const auto& api_hooks = hook::get_or_create();
    const auto& functions = api_hooks->replaceable_tables[std::string(module_name)];
    if (const auto function = functions.find(function_name); function != functions.end())
        return function->second;

    return import_address{};
}

void hook::modify_addresses_in_import_tables(const HMODULE module_address, const operation_type operation) noexcept
{
    unsigned long entry_size{};
    for (auto import_descriptor = reinterpret_cast<PIMAGE_IMPORT_DESCRIPTOR>(::ImageDirectoryEntryToData(module_address, TRUE, IMAGE_DIRECTORY_ENTRY_IMPORT, &entry_size)); import_descriptor && import_descriptor->Name; ++import_descriptor)
    {
        std::string module_name = reinterpret_cast<PSTR>(reinterpret_cast<uint8_t*>(module_address) + import_descriptor->Name);
        std::transform(module_name.begin(), module_name.end(), module_name.begin(), ::tolower);

        if (const auto iterator = replaceable_tables.find(module_name); iterator != replaceable_tables.end())
        {
            const auto iat = reinterpret_cast<PIMAGE_THUNK_DATA>(reinterpret_cast<uint8_t*>(module_address) + import_descriptor->FirstThunk);
            modify_addresses_in_table(operation, iat, iterator->second);
        }
    }
}

void hook::modify_addresses_in_table(const operation_type operation, PIMAGE_THUNK_DATA iat, const function_table& addresses) noexcept
{
    for (; iat->u1.Function; ++iat)
    {
        const auto function = reinterpret_cast<PROC*>(&iat->u1.Function);

        for (const auto& [function_name, address] : addresses)
        {
            const auto source = operation == operation_type::modify ? address.original : address.modified;
            const auto destination = operation == operation_type::modify ? address.modified : address.original;

            if (*function == source)
            {
                DWORD old_protect{};
                if (::VirtualProtect(function, sizeof(function), PAGE_EXECUTE_READWRITE, &old_protect))
                    // can also be used ::WriteProcessMemory(::GetCurrentProcess(), function, &destination, sizeof(destination), NULL))
                    iat->u1.Function = reinterpret_cast<ULONGLONG>(destination);
                else
                    ;
            }
        }
    }
}

hook::import_address hook::get_function_address(const HMODULE module_address, PCSTR function_name)
{
    if (!module_address)
        return import_address{};

    const auto& api_hooks = hook::get_or_create();
    if (const auto name = api_hooks->module_names.find(module_address); name != api_hooks->module_names.end())
        return get_function_address(name->second, function_name);

    return import_address{};
}

hook::function_address hook::get_procedure_address(const HMODULE module_address, PCSTR function_name)
{
    const auto& api_hooks = hook::get_or_create();
    const auto original_get_proc_address = api_hooks->get_function_address(::GetModuleHandleA(kernel_module.c_str()), "GetProcAddress").original;
    const auto procedure_address = reinterpret_cast<FARPROC (*)(HMODULE, PCSTR)>(original_get_proc_address)(module_address, function_name);

    if (!procedure_address)
        return procedure_address;

    if (const auto modified_address = get_function_address(module_address, function_name).modified)
        return modified_address;

    return procedure_address;
}

} // namespace windows
